"""
Decorator functions for registering components with the DIL_Frame registry system.

This module provides decorators that make it easy to register classes with
the appropriate registries while maintaining proper typing information.
These decorators are fundamental to the plugin architecture of DIL_Frame,
allowing for modular extension of system components.
"""

from typing import Callable, Optional, Type, TypeVar

from .registry import (
    LEARNER_REGISTRY,
    TRAINING_STRATEGY_REGISTRY,
    DRIFT_COMPENSATION_REGISTRY,
    PROTOTYPE_MANAGER_REGISTRY,
    MEMORY_MANAGER_REGISTRY,
    CLASSIFIER_ALIGNMENT_REGISTRY,
    EVALUATION_MANAGER_REGISTRY,
)
from learners.interfaces import (
    LearnerInterface,
    MemoryInterface,
    PrototypeInterface,
    ClassifierAlignmentInterface,
    DriftCompensationInterface,
    TrainingInterface,
    EvaluationInterface,
)

T = TypeVar("T")


def register_learner(name: Optional[str] = None) -> Callable[[Type[LearnerInterface]], Type[LearnerInterface]]:
    """Register a class with the learner registry.

    The decorated class must implement LearnerInterface which defines the core
    functionality for incremental learning algorithms. Learners are responsible
    for the main learning process, including model updates, prediction, and
    managing the incremental learning cycle.

    Args:
        name: Optional name to register the class under. If None, uses class name.
              The name should be unique across all registered learners.

    Returns:
        Decorator function that registers the class with the learner registry
        and returns the original class unchanged.

    Example:
        @register_learner("my_learner")
        class MyLearner(LearnerInterface):
            # Implement required methods
            ...
    """
    return LEARNER_REGISTRY.register(name)


def register_training_strategy(
    name: Optional[str] = None,
) -> Callable[[Type[TrainingInterface]], Type[TrainingInterface]]:
    """Register a class with the training strategy registry.

    Training strategies define how models are trained in the incremental learning
    process, including batch size management, epoch counts, optimization methods,
    and loss functions. The decorated class must implement TrainingInterface.

    Args:
        name: Optional name to register the strategy under. If None, uses class name.
              The name should be unique within the training strategy registry.

    Returns:
        Decorator function that registers the class with the training strategy registry
        and returns the original class unchanged.

    Example:
        @register_training_strategy("custom_training")
        class CustomTrainingStrategy(TrainingInterface):
            # Implement required methods
            ...
    """
    return TRAINING_STRATEGY_REGISTRY.register(name)


def register_drift_compensation(
    name: Optional[str] = None,
) -> Callable[[Type[DriftCompensationInterface]], Type[DriftCompensationInterface]]:
    """Register a class with the drift compensation registry.

    Drift compensation components detect and mitigate concept drift in data streams,
    ensuring model adaptability over time. These components are crucial for maintaining
    classification accuracy in non-stationary environments.

    The decorated class must implement DriftCompensationInterface.

    Args:
        name: Optional name to register the component under. If None, uses class name.
              The name should be unique across all drift compensation components.

    Returns:
        Decorator function that registers the class with the drift compensation registry
        and returns the original class unchanged.

    Example:
        @register_drift_compensation("adaptive_drift_handler")
        class AdaptiveDriftHandler(DriftCompensationInterface):
            # Implement required methods
            ...
    """
    return DRIFT_COMPENSATION_REGISTRY.register(name)


def register_prototype_manager(
    name: Optional[str] = None,
) -> Callable[[Type[PrototypeInterface]], Type[PrototypeInterface]]:
    """Register a class with the prototype manager registry.

    Prototype managers handle the creation, update, and retrieval of prototypes
    that represent class distributions in feature space. They play a key role in
    prototype-based incremental learning approaches, helping to maintain compact
    representations of class knowledge.

    The decorated class must implement PrototypeInterface.

    Args:
        name: Optional name to register the manager under. If None, uses class name.
              The name should be unique within the prototype manager registry.

    Returns:
        Decorator function that registers the class with the prototype manager registry
        and returns the original class unchanged.

    Example:
        @register_prototype_manager("mean_prototype_manager")
        class MeanPrototypeManager(PrototypeInterface):
            # Implement required methods
            ...
    """
    return PROTOTYPE_MANAGER_REGISTRY.register(name)


def register_memory_manager(name: Optional[str] = None) -> Callable[[Type[MemoryInterface]], Type[MemoryInterface]]:
    """Register a class with the memory manager registry.

    Memory managers are responsible for storing, retrieving, and maintaining
    examples from previous learning sessions. They are essential for combating
    catastrophic forgetting in incremental learning scenarios by providing
    access to historical data during model updates.

    The decorated class must implement MemoryInterface.

    Args:
        name: Optional name to register the manager under. If None, uses class name.
              The name should be unique within the memory manager registry.

    Returns:
        Decorator function that registers the class with the memory manager registry
        and returns the original class unchanged.

    Example:
        @register_memory_manager("reservoir_memory")
        class ReservoirMemory(MemoryInterface):
            # Implement required methods
            ...
    """
    return MEMORY_MANAGER_REGISTRY.register(name)


def register_classifier_alignment(
    name: Optional[str] = None,
) -> Callable[[Type[ClassifierAlignmentInterface]], Type[ClassifierAlignmentInterface]]:
    """Register a class with the classifier alignment registry.

    Classifier alignment components handle the integration of new classification
    capabilities with existing ones, ensuring consistent prediction behavior
    as the model evolves over time. This is particularly important when adding
    new classes to an existing model.

    The decorated class must implement ClassifierAlignmentInterface.

    Args:
        name: Optional name to register the alignment under. If None, uses class name.
              The name should be unique within the classifier alignment registry.

    Returns:
        Decorator function that registers the class with the classifier alignment registry
        and returns the original class unchanged.

    Example:
        @register_classifier_alignment("weight_alignment")
        class WeightAlignmentStrategy(ClassifierAlignmentInterface):
            # Implement required methods
            ...
    """
    return CLASSIFIER_ALIGNMENT_REGISTRY.register(name)


def register_evaluation_manager(
    name: Optional[str] = None,
) -> Callable[[Type[EvaluationInterface]], Type[EvaluationInterface]]:
    """Register a class with the evaluation manager registry.

    Evaluation managers define protocols for assessing model performance in
    incremental learning scenarios. They handle metrics collection, testing
    procedures, and evaluation schemes that are specifically designed to
    measure performance on both current and previously learned tasks.

    The decorated class must implement EvaluationInterface.

    Args:
        name: Optional name to register the manager under. If None, uses class name.
              The name should be unique within the evaluation manager registry.

    Returns:
        Decorator function that registers the class with the evaluation manager registry
        and returns the original class unchanged.

    Example:
        @register_evaluation_manager("incremental_evaluator")
        class IncrementalEvaluator(EvaluationInterface):
            # Implement required methods
            ...
    """
    return EVALUATION_MANAGER_REGISTRY.register(name)


def register_multi(*decorators: Callable) -> Callable[[Type[T]], Type[T]]:
    """
    Apply multiple registration decorators to a single class.

    This utility allows a class to be registered with multiple registries
    simultaneously, which is useful for components that implement multiple
    interfaces or serve multiple roles in the DIL_Frame architecture.

    Args:
        *decorators: Registration decorators to apply, in the order they should
                    be applied (from innermost to outermost).

    Returns:
        Decorator that applies all provided decorators in sequence to the
        decorated class.

    Example:
        @register_multi(
            register_drift_compensation("hybrid_component"),
            register_prototype_manager("hybrid_component")
        )
        class HybridComponent(DriftCompensationInterface, PrototypeInterface):
            # Implement required methods from both interfaces
            ...
    """

    def multi_decorator(cls: Type[T]) -> Type[T]:
        result = cls
        for decorator in decorators:
            result = decorator(result)
        return result

    return multi_decorator
